import os

def count_objects_in_yolo_labels(folder_path):
    """
    统计YOLO格式标签文件夹中的对象数量
    
    Args:
        folder_path (str): 包含test、train、val子文件夹的根目录路径
    """
    # 定义要统计的子文件夹
    subfolders = ['test', 'train', 'val']
    
    total_objects = 0
    folder_stats = {}
    
    for subfolder in subfolders:
        subfolder_path = os.path.join(folder_path, subfolder)
        
        # 检查子文件夹是否存在
        if not os.path.exists(subfolder_path):
            print(f"警告: 子文件夹 '{subfolder}' 不存在于路径 {folder_path}")
            continue
        
        subfolder_objects = 0
        file_count = 0
        
        # 遍历子文件夹中的所有txt文件
        for filename in os.listdir(subfolder_path):
            if filename.endswith('.txt'):
                file_path = os.path.join(subfolder_path, filename)
                
                try:
                    with open(file_path, 'r', encoding='utf-8') as file:
                        lines = file.readlines()
                        # 统计非空行（每个非空行代表一个对象）
                        object_count = len([line for line in lines if line.strip()])
                        subfolder_objects += object_count
                        file_count += 1
                        
                except Exception as e:
                    print(f"读取文件 {file_path} 时出错: {e}")
        
        folder_stats[subfolder] = {
            'objects': subfolder_objects,
            'files': file_count
        }
        total_objects += subfolder_objects
    
    # 打印统计结果
    print("=" * 50)
    print("YOLO标签文件对象统计结果")
    print("=" * 50)
    
    for subfolder, stats in folder_stats.items():
        print(f"{subfolder}文件夹:")
        print(f"  - 标签文件数量: {stats['files']}")
        print(f"  - 对象总数: {stats['objects']}")
        print()
    
    print(f"所有文件夹对象总数: {total_objects}")
    print("=" * 50)

if __name__ == "__main__":
    # 在这里指定你的文件夹路径
    folder_path = r"F:\论文数据集训练备份\paper_dataset\mydataset\my_usv_aug\integrated_dataset_seaships_usv_weak_nousv5\labels"  # 请替换为你的实际路径
    
    # 检查路径是否存在
    if not os.path.exists(folder_path):
        print(f"错误: 指定的文件夹路径不存在: {folder_path}")
    else:
        count_objects_in_yolo_labels(folder_path)